import cv2 as cv
import numpy as np
import LineTrack as Tracker

black = np.ones((36, 30, 1), dtype=np.uint8)
mask = black.copy()
mask[3:33, 2:28] = 255
kernel = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))


def Process(img):
	resized = cv.resize(img, (30, 36))
	gray = cv.cvtColor(resized, cv.COLOR_BGR2GRAY)
	_, thresh = cv.threshold(gray, 100, 255, cv.THRESH_BINARY_INV)
	res = cv.bitwise_and(thresh, mask)  # , mask=mask)
	contours, _ = cv.findContours(res, cv.INTER_LINEAR, cv.CHAIN_APPROX_NONE)
	for contour in contours:
		if cv.contourArea(contour) > 40:
			cv.drawContours(black.copy(), contour, -1, (255, 0, 0), -1)
			close = cv.morphologyEx(black.copy(), cv.MORPH_CLOSE, kernel, iterations=2)
			return cv.resize(close, (40, 40))
		else:
			return cv.resize(black, (40, 40))


class Director:
	svm = cv.ml.SVM_load('resources/svm_data.dat')
	__target = 3
	direction = Tracker.Straight

	def SetTarget(self, img):
		if self.__target == 0:
			for i in range(210, 430, 10):
				for j in range(100, 200, 5):
					# print(i, j)
					f = img[i:i+100, j:j+100]
					f = Process(f)
					f = f.astype(np.float32).reshape(-1, 64)
					num_pre = self.svm.predict(f)[1][0][0]
					self.__target = num_pre

	def Direct(self, img_list, img):
		flag = 0
		if len(img_list):
			if self.__target:
				for i in range(len(img_list)):
					img = img_list[i]
					img = Process(img)
					img = img.astype(np.float32).reshape(-1, 64)
					num_pre = self.svm.predict(img)[1][0][0]
					print(num_pre)
					if int(num_pre) == self.__target:
						if i < 2:
							self.direction = Tracker.Left
						else:
							self.direction = Tracker.Right
						flag = 1
					else:
						self.direction = Tracker.Straight
		else:
			self.SetTarget(img)

		return flag



